import os, sys
import json
import sqlite3
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pathlib import Path
import pymysql

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../')))
from insight_agent.utils.http_post import get_insight_request, inform_insight_thinking, update_insight_result

# 加载环境变量
load_dotenv(override=True)

# 创建MCP
mcp = FastMCP("AgentSQLiteServer")

# MySQL连接配置
MYSQL_CONFIG = {
    'host': os.getenv('MYSQL_HOST', 'localhost'),
    'port': int(os.getenv('MYSQL_PORT', 3306)),
    'user': os.getenv('MYSQL_USER', 'root'),
    'password': os.getenv('MYSQL_PASSWORD', ''),
    'database': os.getenv('MYSQL_DATABASE', 'user_insight'),
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor  # 使结果自动转为字典格式
}

# 构建文件读写路径
file_folder_path = Path(__file__).resolve().parent.parent.parent / "files"

# 数据读取模式，例如 SQLITE, MYSQL等
data_connection_mode = "SQLITE"

# Tools
@mcp.tool()
async def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用于在指定SQLite数据库上运行一段SQL代码，完成数据查询相关工作，
    并且当前函数是使用Python标准库sqlite3连接SQLite数据库。
    本函数只负责运行SQL代码并进行数据查询。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对SQLite数据库进行查询，并获得相关信息
    :return：sql_query在SQLite中的运行结果。
    """
    inform_insight_thinking("agent_time_series", f'tools - sql_inquire')
    if data_connection_mode == "SQLITE":
        # 创建连接（SQLite使用文件路径，若文件不存在会自动创建）
        connection = sqlite3.connect(file_folder_path / 'sqlite_data.db')
        # 设置行工厂，让结果包含字段名
        connection.row_factory = sqlite3.Row

        try:
            cursor = connection.cursor()
            cursor.execute(sql_query)
            # 获取所有查询结果
            results = cursor.fetchall()
            # 转换为包含字段名的字典列表
            formatted_results = [dict(row) for row in results]
        except Exception as e:
            # 捕获并返回错误信息
            return json.dumps({"error": str(e)}, ensure_ascii=False)
        finally:
            connection.close()

        # 将结果以JSON字符串形式返回
        return json.dumps(formatted_results, ensure_ascii=False)

    elif data_connection_mode == "MYSQL":
        connection = None
        try:
            # 创建MySQL连接
            connection = pymysql.connect(**MYSQL_CONFIG)

            with connection.cursor() as cursor:
                cursor.execute(sql_query)

                # 获取查询结果
                if sql_query.strip().upper().startswith('SELECT'):
                    results = cursor.fetchall()
                    return json.dumps(results, ensure_ascii=False)
                else:
                    # 对于非SELECT语句，返回影响的行数
                    connection.commit()
                    return json.dumps({"affected_rows": cursor.rowcount}, ensure_ascii=False)

        except Exception as e:
            # 回滚事务并返回错误信息
            if connection:
                connection.rollback()
            return json.dumps({"error": str(e)}, ensure_ascii=False)
        finally:
            if connection:
                connection.close()
    else:
        raise Exception("Unsupported Data Connection Mode.")

if __name__ == "__main__":
    mcp.run(transport="stdio")